"""
数据库配置和会话管理
提供SQLAlchemy数据库连接和会话管理功能
"""

import logging
from sqlalchemy import create_engine, event
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker
from sqlalchemy.orm import sessionmaker, Session
from sqlalchemy.pool import QueuePool
from typing import Generator, AsyncGenerator
from contextlib import contextmanager, asynccontextmanager

from .config import settings
from ..models.base import Base

logger = logging.getLogger(__name__)

# 数据库引擎
engine = None
async_engine = None
SessionLocal = None
AsyncSessionLocal = None


def create_database_engine():
    """
    创建同步数据库引擎
    配置连接池和优化参数
    """
    global engine
    
    if engine is None:
        # 根据数据库URL类型选择不同的配置
        if "sqlite" in settings.database_url:
            # SQLite配置
            engine = create_engine(
                settings.database_url.replace("+aiomysql", "+pymysql"),
                connect_args={"check_same_thread": False},  # SQLite特有配置
                echo=False,  # 关闭SQL日志输出
                pool_pre_ping=True,   # 连接前检查
            )
        elif "mysql" in settings.database_url:
            # MySQL配置
            engine = create_engine(
                settings.database_url.replace("+aiomysql", "+pymysql"),
                poolclass=QueuePool,
                pool_size=10,         # 连接池大小
                max_overflow=20,      # 最大溢出连接数
                pool_pre_ping=True,   # 连接前检查
                pool_recycle=3600,    # 连接回收时间（1小时）
                echo=False,  # 关闭SQL日志输出
                connect_args={
                    "charset": "utf8mb4",
                    "autocommit": True
                }
            )
        else:
            # PostgreSQL配置
            engine = create_engine(
                settings.database_url,
                poolclass=QueuePool,
                pool_size=10,         # 连接池大小
                max_overflow=20,      # 最大溢出连接数
                pool_pre_ping=True,   # 连接前检查
                pool_recycle=3600,    # 连接回收时间（1小时）
                echo=False,  # 关闭SQL日志输出
            )
        
        # 添加连接事件监听器
        @event.listens_for(engine, "connect")
        def set_sqlite_pragma(dbapi_connection, connection_record):
            """为SQLite设置优化参数"""
            if "sqlite" in settings.database_url:
                cursor = dbapi_connection.cursor()
                # 启用外键约束
                cursor.execute("PRAGMA foreign_keys=ON")
                # 设置WAL模式提高并发性能
                cursor.execute("PRAGMA journal_mode=WAL")
                # 优化同步设置
                cursor.execute("PRAGMA synchronous=NORMAL")
                cursor.close()
        
        # 数据库引擎初始化完成（日志已简化）
    
    return engine


async def create_async_database_engine():
    """
    创建异步数据库引擎
    配置连接池和优化参数
    """
    global async_engine
    
    if async_engine is None:
        # 使用异步引擎
        async_engine = create_async_engine(
            settings.database_url,
            echo=False,  # 关闭SQL日志输出
            pool_size=10,
            max_overflow=20,
            pool_pre_ping=True,
            pool_recycle=3600,
        )
        # 异步数据库引擎初始化完成（日志已简化）
    
    return async_engine


def create_session_factory():
    """创建同步会话工厂"""
    global SessionLocal
    
    if SessionLocal is None:
        engine = create_database_engine()
        SessionLocal = sessionmaker(
            autocommit=False,
            autoflush=False,
            bind=engine
        )
        # 同步数据库会话工厂已创建（日志已简化）
    
    return SessionLocal


async def create_async_session_factory():
    """创建异步会话工厂"""
    global AsyncSessionLocal
    
    if AsyncSessionLocal is None:
        engine = await create_async_database_engine()
        AsyncSessionLocal = async_sessionmaker(
            autocommit=False,
            autoflush=False,
            bind=engine,
            class_=AsyncSession
        )
        # 异步数据库会话工厂已创建（日志已简化）
    
    return AsyncSessionLocal


def get_database_session() -> Generator[Session, None, None]:
    """
    获取同步数据库会话
    用于FastAPI依赖注入
    
    Yields:
        Session: 数据库会话对象
    """
    session_factory = create_session_factory()
    session = session_factory()
    try:
        yield session
    except Exception as e:
        logger.error(f"数据库会话错误: {e}")
        session.rollback()
        raise
    finally:
        session.close()


async def get_async_database_session() -> AsyncGenerator[AsyncSession, None]:
    """
    获取异步数据库会话
    用于FastAPI依赖注入
    
    Yields:
        AsyncSession: 异步数据库会话对象
    """
    session_factory = await create_async_session_factory()
    session = session_factory()
    try:
        yield session
    except Exception as e:
        logger.error(f"异步数据库会话错误: {e}")
        await session.rollback()
        raise
    finally:
        await session.close()


@contextmanager
def get_db_session():
    """
    获取同步数据库会话上下文管理器
    用于服务层直接使用
    
    Yields:
        Session: 数据库会话对象
    """
    session_factory = create_session_factory()
    session = session_factory()
    try:
        yield session
        session.commit()
    except Exception as e:
        logger.error(f"数据库操作失败: {e}")
        session.rollback()
        raise
    finally:
        session.close()


@asynccontextmanager
async def get_async_db_session():
    """
    获取异步数据库会话上下文管理器
    用于服务层直接使用
    
    Yields:
        AsyncSession: 异步数据库会话对象
    """
    session_factory = await create_async_session_factory()
    session = session_factory()
    try:
        yield session
        await session.commit()
    except Exception as e:
        logger.error(f"异步数据库操作失败: {e}")
        await session.rollback()
        raise
    finally:
        await session.close()


def init_database():
    """
    初始化数据库
    创建所有表结构
    """
    try:
        # 导入所有模型以确保它们被注册
        from ..models.visitor import Visitor
        from ..models.ai_person import AiPerson  
        from ..models.image_generation_task import ImageGenerationTask
        
        engine = create_database_engine()
        
        # 创建所有表
        Base.metadata.create_all(bind=engine)
        # 数据库表结构已创建/更新（日志已简化）
        
        return True
    except Exception as e:
        logger.error(f"数据库初始化失败: {e}")
        return False


def check_database_connection() -> bool:
    """
    检查数据库连接状态
    
    Returns:
        bool: 连接是否正常
    """
    try:
        engine = create_database_engine()
        with engine.connect() as connection:
            connection.execute("SELECT 1")
        # 数据库连接检查通过（日志已简化）
        return True
    except Exception as e:
        logger.error(f"数据库连接检查失败: {e}")
        return False


async def close_database():
    """关闭数据库连接"""
    global engine, async_engine
    if engine:
        engine.dispose()
        engine = None
    if async_engine:
        await async_engine.dispose()
        async_engine = None
    # 数据库连接已关闭（日志已简化）


# 数据库会话依赖
DatabaseSession = get_database_session
AsyncDatabaseSession = get_async_database_session 